{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IERG 5350 Assignment 1: Tabular Reinforcement Learning\n",
    "\n",
    "*2021-2022 1st term, IERG 5350: Reinforcement Learning. Department of Information Engineering, The Chinese University of Hong Kong. Course Instructor: Professor ZHOU Bolei. Assignment author: PENG Zhenghao.*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Student Name | Student ID |\n",
    "| :----: | :----: |\n",
    "| TYPE_YOUR_NAME_HERE | TYPE_YOUR_STUDENT_ID_HERE |\n",
    "\n",
    "------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the assignment 1 of our RL course. The objective of this assignment is for you to understand the classic methods used in tabular reinforcement learning. \n",
    "\n",
    "This assignment has the following sections:\n",
    "\n",
    " - Section 1: Warm-up on the RL environment (35 points)\n",
    " - Section 2: Implementation of the model-based family of algorithms: policy iteration and value iteration. (65 points)\n",
    "\n",
    "You need to go through this self-contained notebook, which contains **21 TODOs** in part of the cells and has special `[TODO]` signs. You need to finish all TODOs.\n",
    "\n",
    "You are encouraged to add more code on extra cells at the end of the each section to investigate the problems you think interesting. At the end of the file, we leave a place for you to write comments optionaly (Yes, please give us either negative or positive rewards so we can keep improving the assignment!). \n",
    "\n",
    "Please report any code bugs to us via **github issues**.\n",
    "\n",
    "Before you get start, remember to follow the instruction at \n",
    "https://github.com/cuhkrlcourse/ierg5350-assignment-2021/tree/main/assignment0\n",
    "to setup your environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now start running the cells sequentially (by `ctrl + enter` or `shift + enter`) to avoid unnecessary errors by skipping some cells. \n",
    "\n",
    "\n",
    "## Section 1: Warm-up on the RL environment\n",
    "\n",
    "(35/100 points)\n",
    "\n",
    "In this section, we will go through the basic concepts of RL environments using OpenAI Gym. Besides, you will get the first sense of the toy environment we will use in the rest of the assignment.\n",
    "\n",
    "Every Gym environment should contain the following attributes:\n",
    "\n",
    "1. `env.step(action)` To step the environment by applying `action`. Will return four things: `observation, reward, done, info`, wherein `done` is a boolean value indicating whether this **episode** is finished. `info` may contain some information the user is interested in, we do not use it.\n",
    "2. `env.reset()` To reset the environment, back to the initial state. Will return the initial observation.\n",
    "3. `env.render()` To render the current state of the environment for human-being\n",
    "4. `env.action_space` The allowed action format. In our case, it is `Discrete(4)` which means the action is an integer in the range [0, 1, 2, 3]. Therefore the `action` for `step(action)` should obey the limit of the action space.\n",
    "5. `env.observation_space` The observation space.\n",
    "6. `env.seed(seed)` To set the random seed of the environment. So the result is replicable.\n",
    "\n",
    "Note that the word **episode** means the process that an agent interacts with the environment from the initial state to the terminal state. Within one episode, the agent will only receive one `done=True`, when it goes to the terminal state (the agent is dead or the game is over).\n",
    "\n",
    "We will use \"FrozenLake8x8-v1\" as our environment. In this environment, the agent controls the movement of a character in a grid world. Some tiles of the grid are walkable, and others lead to the agent falling into the water. Additionally, the movement direction of the agent is uncertain and only partially depends on the chosen direction. The agent is rewarded for finding a walkable path to a goal tile. The meaning of each character:\n",
    "\n",
    "1. S : starting point, safe\n",
    "2. F : frozen surface, safe\n",
    "3. H : hole, fall to your doom\n",
    "4. G : goal, where the frisbee is located\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# Import some packages that we need to use\n",
    "from utils import *\n",
    "import gym\n",
    "import numpy as np\n",
    "from collections import deque"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 1.1: Make the environment\n",
    "\n",
    "You need to know \n",
    "\n",
    "1. How to make an environment\n",
    "2. How to set the random seed of environment\n",
    "3. What is observation space and action space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current observation space: Discrete(64)\n",
      "Current action space: Discrete(4)\n",
      "0 in action space? True\n",
      "5 in action space? False\n"
     ]
    }
   ],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# [TODO] Just a reminder. Do you add your name and student \n",
    "# ID in the table at top of the notebook?\n",
    "pass\n",
    "\n",
    "# Create the environment\n",
    "env = gym.make('FrozenLake8x8-v1')\n",
    "\n",
    "# You need to reset the environment immediately after instantiating env. \n",
    "env.reset()  # [TODO] uncomment this line\n",
    "# Seed the environment\n",
    "env.seed(0)  # [TODO] uncomment this line\n",
    "\n",
    "print(\"Current observation space: {}\".format(env.observation_space))\n",
    "print(\"Current action space: {}\".format(env.action_space))\n",
    "print(\"0 in action space? {}\".format(env.action_space.contains(0)))\n",
    "print(\"5 in action space? {}\".format(env.action_space.contains(5)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 1.2: Play the environment with random actions\n",
    "\n",
    "You need to know \n",
    "\n",
    "1. How to step the environment\n",
    "2. How to render the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Right)\n",
      "SFFFFFFF\n",
      "FFFFFFFF\n",
      "FFF\u001B[41mH\u001B[0mFFFF\n",
      "FFFFFHFF\n",
      "FFFHFFFF\n",
      "FHHFFFHF\n",
      "FHFFHFHF\n",
      "FFFHFFFG\n",
      "Current observation: 19\n",
      "Current reward: 0.0\n",
      "Whether we are done: True\n",
      "info: {'prob': 0.3333333333333333}\n"
     ]
    }
   ],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# Run 1000 steps for test, terminate if done.\n",
    "# You can run this cell multiples times.\n",
    "env.reset()\n",
    "\n",
    "while True:\n",
    "    # take random action\n",
    "    # [TODO] Uncomment next line\n",
    "    obs, reward, done, info = env.step(env.action_space.sample())\n",
    "\n",
    "    # render the environment\n",
    "    env.render()  # [TODO] Uncomment this line\n",
    "\n",
    "    print(\"Current observation: {}\\nCurrent reward: {}\\n\"\n",
    "          \"Whether we are done: {}\\ninfo: {}\".format(\n",
    "        obs, reward, done, info\n",
    "    ))\n",
    "    wait(sleep=0.5)\n",
    "\n",
    "    # [TODO] terminate the loop if done\n",
    "    if done:\n",
    "        env.reset()\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 1.3: Define the evaluation function to value the random baseline\n",
    "\n",
    "Now we need to define an evaluation function to evaluate a given policy (a function where the input is observation and the output is action). This is convenient for future evaluation.\n",
    "\n",
    "As a reminder, you should create a `FrozenLake8x8-v1` environment instance by default, reset it after each episode (and at the beginning), step the environment, and terminate episode if done.\n",
    "\n",
    "After implementing the `evaluate` function, run the next cell to check whether you are right."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "def _render_helper(env):\n",
    "    env.render()\n",
    "    wait(sleep=0.2)\n",
    "\n",
    "\n",
    "def evaluate(policy, num_episodes, seed=0, env_name='FrozenLake8x8-v1', render=False):\n",
    "    \"\"\"[TODO] You need to implement this function by yourself. It\n",
    "    evaluate the given policy and return the mean episode reward.\n",
    "    We use `seed` argument for testing purpose.\n",
    "    You should pass the tests in the next cell.\n",
    "\n",
    "    :param policy: a function whose input is an interger (observation)\n",
    "    :param num_episodes: number of episodes you wish to run\n",
    "    :param seed: an interger, used for testing.\n",
    "    :param env_name: the name of the environment\n",
    "    :param render: a boolean flag. If true, please call _render_helper\n",
    "    function.\n",
    "    :return: the averaged episode reward of the given policy.\n",
    "    \"\"\"\n",
    "\n",
    "    # Create environment (according to env_name, we will use env other than 'FrozenLake8x8-v0')\n",
    "    env = gym.make(env_name)\n",
    "\n",
    "    # Seed the environment\n",
    "    env.seed(seed)\n",
    "\n",
    "    # Build inner loop to run.\n",
    "    # For each episode, do not set the limit.\n",
    "    # Only terminate episode (reset environment) when done = True.\n",
    "    # The episode reward is the sum of all rewards happen within one episode.\n",
    "    # Call the helper function `render(env)` to render\n",
    "    rewards = []\n",
    "    for i in range(num_episodes):\n",
    "        # reset the environment\n",
    "        obs = env.reset()\n",
    "        act = policy(obs)\n",
    "        \n",
    "        ep_reward = 0\n",
    "        while True:\n",
    "            # [TODO] run the environment and terminate it if done, collect the\n",
    "            # reward at each step and sum them to the episode reward.\n",
    "            obs, reward, done, info = env.step(act)\n",
    "            ep_reward += reward\n",
    "            act = policy(obs)\n",
    "            if done:\n",
    "                break\n",
    "        \n",
    "        rewards.append(ep_reward)\n",
    "\n",
    "    return np.mean(rewards)\n",
    "\n",
    "# [TODO] Run next cell to test your implementation!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Passed!\n",
      "\n",
      "As a baseline, the mean episode reward of a hand-craft agent is:  0.065\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# Run this cell to test the correctness of your implementation of `evaluate`.\n",
    "LEFT = 0\n",
    "DOWN = 1\n",
    "RIGHT = 2\n",
    "UP = 3\n",
    "\n",
    "def expert(obs):\n",
    "    \"\"\"Go down if agent at the right edge, otherwise go right.\"\"\"\n",
    "    return DOWN if (obs + 1) % 8 == 0 else RIGHT\n",
    "\n",
    "def assert_equal(seed, value, env_name):\n",
    "    ret = evaluate(expert, 1000, seed, env_name=env_name)\n",
    "    assert ret == value, \\\n",
    "    \"When evaluate on seed {}, 1000 episodes, in {} environment, the \" \\\n",
    "    \"averaged reward should be {}. But you get {}.\" \\\n",
    "    \"\".format(seed, env_name, value, ret)\n",
    "      \n",
    "assert_equal(0, 0.065, 'FrozenLake8x8-v1')\n",
    "assert_equal(1, 0.059, 'FrozenLake8x8-v1')\n",
    "assert_equal(2, 0.055, 'FrozenLake8x8-v1')\n",
    "\n",
    "assert_equal(0, 0.026, 'FrozenLake-v1')\n",
    "assert_equal(1, 0.034, 'FrozenLake-v1')\n",
    "assert_equal(2, 0.028, 'FrozenLake-v1')\n",
    "\n",
    "print(\"Test Passed!\")\n",
    "print(\"\\nAs a baseline, the mean episode reward of a hand-craft \"\n",
    "      \"agent is: \", evaluate(expert, 1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congraduation! You have finished section 1 (if and only if not error happen at the above codes).\n",
    "\n",
    "If you want to do more investigation, feel free to open new cells via `Esc + B` after the next cells and write codes in it, so that you can reuse some result in this notebook. Remember to write sufficient comments and documents to let others know what you are doing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# You can do more inverstigation here if you wish. Leave it blank if you don't.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------\n",
    "\n",
    "## Section 2: Model-based Tabular RL\n",
    "\n",
    "(65/100 points)\n",
    "\n",
    "We have learned how to use the Gym environment to run an episode, as well as how to interact between the agent (policy) and environment via `env.step(action)` to collect observation, reward, done, and possible extra information.\n",
    "\n",
    "Now we need to build the basic tabular RL algorithm to solve this environment. **Note that compared to the model-free methods in the Sec.3, the algorithms in this section needs to access the internal information of the environment, namely the transition dynamics**. In our case, given a state and an action, we need to know which state current environment would jump to, and the probability of this happens, and the reward if the transition happens. You will see that we provide you a helper function `trainer._get_transitions(state, action)` that takes state and action as input and return you a list of possible transitions.\n",
    "\n",
    "You will use a class to represent a Trainer, which seems to be over-complex for tabular RL. But we will use the same framework in the future assignments, or even in your future research. So it would be helpful for you to get familiar with how to implement an RL algorithm in a class-orientetd programming style, as a first step toward the implementation of state of the art RL algorithm in the future."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "class TabularRLTrainerAbstract:\n",
    "    \"\"\"This is the abstract class for tabular RL trainer. We will inherent the specify \n",
    "    algorithm's trainer from this abstract class, so that we can reuse the codes like\n",
    "    getting the dynamic of the environment (self._get_transitions()) or rendering the\n",
    "    learned policy (self.render()).\"\"\"\n",
    "    \n",
    "    def __init__(self, env_name='FrozenLake8x8-v1', model_based=True):\n",
    "        self.env_name = env_name\n",
    "        self.env = gym.make(self.env_name)\n",
    "        self.action_dim = self.env.action_space.n\n",
    "        self.obs_dim = self.env.observation_space.n\n",
    "        \n",
    "        self.model_based = model_based\n",
    "\n",
    "    def _get_transitions(self, state, act):\n",
    "        \"\"\"Query the environment to get the transition probability,\n",
    "        reward, the next state, and done given a pair of state and action.\n",
    "        We implement this function for you. But you need to know the \n",
    "        return format of this function.\n",
    "        \"\"\"\n",
    "        self._check_env_name()\n",
    "        assert self.model_based, \"You should not use _get_transitions in \" \\\n",
    "            \"model-free algorithm!\"\n",
    "        \n",
    "        # call the internal attribute of the environments.\n",
    "        # `transitions` is a list contain all possible next states and the \n",
    "        # probability, reward, and termination indicater corresponding to it\n",
    "        transitions = self.env.env.P[state][act]\n",
    "\n",
    "        # Given a certain state and action pair, it is possible\n",
    "        # to find there exist multiple transitions, since the \n",
    "        # environment is not deterministic.\n",
    "        # You need to know the return format of this function: a list of dicts\n",
    "        ret = []\n",
    "        for prob, next_state, reward, done in transitions:\n",
    "            ret.append({\n",
    "                \"prob\": prob,\n",
    "                \"next_state\": next_state,\n",
    "                \"reward\": reward,\n",
    "                \"done\": done\n",
    "            })\n",
    "        return ret\n",
    "    \n",
    "    def _check_env_name(self):\n",
    "        assert self.env_name.startswith('FrozenLake')\n",
    "\n",
    "    def print_table(self):\n",
    "        \"\"\"print beautiful table, only work for FrozenLake8X8-v0 env. We \n",
    "        write this function for you.\"\"\"\n",
    "        self._check_env_name()\n",
    "        print_table(self.table)\n",
    "\n",
    "    def train(self):\n",
    "        \"\"\"Conduct one iteration of learning.\"\"\"\n",
    "        raise NotImplementedError(\"You need to override the \"\n",
    "                                  \"Trainer.train() function.\")\n",
    "\n",
    "    def evaluate(self):\n",
    "        \"\"\"Use the function you write to evaluate current policy.\n",
    "        Return the mean episode reward of 1000 episodes when seed=0.\"\"\"\n",
    "        result = evaluate(self.policy, 1000, env_name=self.env_name)\n",
    "        return result\n",
    "\n",
    "    def render(self):\n",
    "        \"\"\"Reuse your evaluate function, render current policy \n",
    "        for one episode when seed=0\"\"\"\n",
    "        evaluate(self.policy, 1, render=True, env_name=self.env_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2.1: Policy Iteration\n",
    "\n",
    "Recall the idea of policy iteration: \n",
    "\n",
    "1. Update the state value function, given all possible transitions in the environment.\n",
    "2. Find the best policy that earns highest value under current state value function.\n",
    "3. If the best policy is identical to the previous one then stop the training. Otherwise, return to step 1.\n",
    "\n",
    "In step 1, the way to update the state value function is by \n",
    "\n",
    "$$v_{k+1} = E_{s'}[r(s, a)+\\gamma v_{k}(s')]$$\n",
    "\n",
    "wherein the $a$ is given by current policy, $s'$ is next state, $r$ is the reward, $v_{k}(s')$ is the next state value given by the old (not updated yet) value function. The expectation is computed among all possible transitions (given a state and action pair, it is possible to have many different next states, since the environment is not deterministic).\n",
    "\n",
    "In step 2, the best policy is the one that takes the action with maximal expected return given a state:\n",
    "\n",
    "$$a = {argmax}_a E_{s'}[r(s, a) + \\gamma v_{k}(s')]$$\n",
    "\n",
    "Policy iteration algorithm has an outer loop (update policy, step 1 to 3) and an inner loop (fit the value function, within step 1). \n",
    "\n",
    "In each outer loop, we call once `trainer.train()`, where we call `trainer.update_value_function()` once to update the value function (the state value table). \n",
    "\n",
    "After that we call `trainer.update_policy()` to update the current policy. \n",
    "\n",
    "`trainer` object has a `trainer.policy` attribute, which is a function that takes observation as input and returns an action.\n",
    "\n",
    "You should implement the trainer following the framework we already wrote for you. Please carefully go through the codes and finish all `TODO` in it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "class PolicyItertaionTrainer(TabularRLTrainerAbstract):\n",
    "    def __init__(self, gamma=1.0, eps=1e-10, env_name='FrozenLake8x8-v1'):\n",
    "        super(PolicyItertaionTrainer, self).__init__(env_name)\n",
    "\n",
    "        # discount factor\n",
    "        self.gamma = gamma\n",
    "\n",
    "        # value function convergence criterion\n",
    "        self.eps = eps\n",
    "\n",
    "        # build the value table for each possible observation\n",
    "        self.table = np.zeros((self.obs_dim,))\n",
    "\n",
    "        # [TODO] you need to implement a random policy at the beginning.\n",
    "        # It is a function that take an integer (state or say observation)\n",
    "        # as input and return an interger (action).\n",
    "        # remember, you can use self.action_dim to get the dimension (range)\n",
    "        # of the action, which is an integer in range\n",
    "        # [0, ..., self.action_dim - 1]\n",
    "        # hint: generating random action at each call of policy may lead to\n",
    "        #  failure of convergence, try generate random actions at initializtion\n",
    "        #  and fix it during the training.\n",
    "        self.policy = lambda _: DOWN if (obs + 1) % 8 == 0 or (obs + 8) < 64 else RIGHT\n",
    "        # test your random policy\n",
    "        test_random_policy(self.policy, self.env)\n",
    "\n",
    "    def train(self):\n",
    "        \"\"\"Conduct one iteration of learning.\"\"\"\n",
    "        # [TODO] value function may be need to be reset to zeros.\n",
    "        # if you think it should, than do it. If not, then move on.\n",
    "        # hint: the value function is equivalent to self.table,\n",
    "        #  a numpy array with length 64.\n",
    "        # self.table = np.zeros((self.obs_dim,))\n",
    "\n",
    "        self.update_value_function()\n",
    "        self.update_policy()\n",
    "\n",
    "    def update_value_function(self):\n",
    "        count = 0  # count the steps of value updates\n",
    "        while True:\n",
    "            old_table = self.table.copy()\n",
    "\n",
    "            for state in range(self.obs_dim):\n",
    "                act = self.policy(state)\n",
    "                transition_list = self._get_transitions(state, act)\n",
    "                \n",
    "                state_value = 0\n",
    "                for transition in transition_list:\n",
    "                    prob = transition['prob']\n",
    "                    reward = transition['reward']\n",
    "                    next_state = transition['next_state']\n",
    "                    done = transition['done']\n",
    "                    \n",
    "                    # [TODO] what is the right state value?\n",
    "                    # hint: you should use reward, self.gamma, old_table, prob,\n",
    "                    # and next_state to compute the state value\n",
    "                    state_value += prob * (reward + self.gamma * old_table[next_state])\n",
    "\n",
    "                # update the state value\n",
    "                self.table[state] = state_value\n",
    "\n",
    "            # [TODO] Compare the old_table and current table to\n",
    "            #  decide whether to break the value update process.\n",
    "            # hint: you should use self.eps, old_table and self.table\n",
    "            should_break = np.sum(np.abs(old_table - self.table)) < self.eps\n",
    "\n",
    "            if should_break:\n",
    "                break\n",
    "            count += 1\n",
    "            if count % 200 == 0:\n",
    "                # disable this part if you think debug message annoying.\n",
    "                print(\"[DEBUG]\\tUpdated values for {} steps. \"\n",
    "                      \"Difference between new and old table is: {}\".format(\n",
    "                    count, np.sum(np.abs(old_table - self.table))\n",
    "                ))\n",
    "            if count > 4000:\n",
    "                print(\"[HINT] Are you sure your codes is OK? It shouldn't be \"\n",
    "                      \"so hard to update the value function. You already \"\n",
    "                      \"use {} steps to update value function within \"\n",
    "                      \"single iteration.\".format(count))\n",
    "            if count > 6000:\n",
    "                raise ValueError(\"Clearly your code has problem. Check it!\")\n",
    "\n",
    "    def update_policy(self):\n",
    "        \"\"\"You need to define a new policy function, given current\n",
    "        value function. The best action for a given state is the one that\n",
    "        has greatest expected return.\n",
    "\n",
    "        To optimize computing efficiency, we introduce a policy table,\n",
    "        which take state as index and return the action given a state.\n",
    "        \"\"\"\n",
    "        policy_table = np.zeros([self.obs_dim, ], dtype=np.int64)\n",
    "\n",
    "        for state in range(self.obs_dim):\n",
    "            state_action_values = [0] * self.action_dim\n",
    "\n",
    "            # [TODO] assign the action with greatest \"value\"\n",
    "            # to policy_table[state]\n",
    "            # hint: what is the proper \"value\" here?\n",
    "            #  you should use table, gamma, reward, prob,\n",
    "            #  next_state and self._get_transitions() function\n",
    "            #  as what we done at self.update_value_function()\n",
    "            #  Bellman equation may help.\n",
    "            best_action = None\n",
    "            for action in range(self.action_dim):\n",
    "                transition_list = self._get_transitions(state, action)\n",
    "                for transition in transition_list:\n",
    "                    prob = transition['prob']\n",
    "                    reward = transition['reward']\n",
    "                    next_state = transition['next_state']\n",
    "                    done = transition['done']\n",
    "                    state_action_values[action] += prob * (reward + self.gamma * self.table[next_state])\n",
    "                best_action = np.argmax(np.array(state_action_values))\n",
    "            \n",
    "            policy_table[state] = best_action\n",
    "\n",
    "        self.policy = lambda obs: policy_table[obs]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have built the Trainer class for policy iteration algorithm. In the following few cells, we will train the agent to solve the problem and evaluate its performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# Managing configurations of your experiments is important for your research.\n",
    "default_pi_config = dict(\n",
    "    max_iteration=1000,\n",
    "    evaluate_interval=1,\n",
    "    gamma=1.0,\n",
    "    eps=1e-10\n",
    ")\n",
    "\n",
    "\n",
    "def policy_iteration(train_config=None):\n",
    "    config = default_pi_config.copy()\n",
    "    if train_config is not None:\n",
    "        config.update(train_config)\n",
    "        \n",
    "    trainer = PolicyItertaionTrainer(gamma=config['gamma'], eps=config['eps'])\n",
    "\n",
    "    old_policy_result = {\n",
    "        obs: -1 for obs in range(trainer.obs_dim)\n",
    "    }\n",
    "\n",
    "    for i in range(config['max_iteration']):\n",
    "        # train the agent\n",
    "        trainer.train()  # [TODO] please uncomment this line\n",
    "\n",
    "        # [TODO] compare the new policy with old policy to check whether\n",
    "        #  should we stop. If new and old policy have same output given any\n",
    "        #  observation, them we consider the algorithm is converged and\n",
    "        #  should be stopped.\n",
    "        new_policy_result = {\n",
    "            state: trainer.policy(state) for state in range(trainer.obs_dim)\n",
    "        }\n",
    "        should_stop = (new_policy_result == old_policy_result)\n",
    "\n",
    "        if should_stop:\n",
    "            print(\"We found policy is not changed anymore at \"\n",
    "                  \"itertaion {}. Current mean episode reward \"\n",
    "                  \"is {}. Stop training.\".format(i, trainer.evaluate()))\n",
    "            break\n",
    "        old_policy_result = new_policy_result\n",
    "\n",
    "        # evaluate the result\n",
    "        if i % config['evaluate_interval'] == 0:\n",
    "            print(\n",
    "                \"[INFO]\\tIn {} iteration, current mean episode reward is {}.\"\n",
    "                \"\".format(i, trainer.evaluate()))\n",
    "\n",
    "            if i > 20:\n",
    "                print(\"You sure your codes is OK? It shouldn't take so many \"\n",
    "                      \"({}) iterations to train a policy iteration \"\n",
    "                      \"agent.\".format(i))\n",
    "\n",
    "    assert trainer.evaluate() > 0.8, \\\n",
    "        \"We expect to get the mean episode reward greater than 0.8. \" \\\n",
    "        \"But you get: {}. Please check your codes.\".format(trainer.evaluate())\n",
    "\n",
    "    return trainer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]\tIn 0 iteration, current mean episode reward is 0.822.\n",
      "[DEBUG]\tUpdated values for 200 steps. Difference between new and old table is: 0.041664161897299004\n",
      "[DEBUG]\tUpdated values for 400 steps. Difference between new and old table is: 0.0022292041480653085\n",
      "[DEBUG]\tUpdated values for 600 steps. Difference between new and old table is: 0.0001184338151329345\n",
      "[DEBUG]\tUpdated values for 800 steps. Difference between new and old table is: 6.291939822350434e-06\n",
      "[DEBUG]\tUpdated values for 1000 steps. Difference between new and old table is: 3.3426684917237104e-07\n",
      "[DEBUG]\tUpdated values for 1200 steps. Difference between new and old table is: 1.7758327711114852e-08\n",
      "[DEBUG]\tUpdated values for 1400 steps. Difference between new and old table is: 9.434331232904825e-10\n",
      "[INFO]\tIn 1 iteration, current mean episode reward is 0.804.\n",
      "[DEBUG]\tUpdated values for 200 steps. Difference between new and old table is: 0.0005348034918033623\n",
      "[DEBUG]\tUpdated values for 400 steps. Difference between new and old table is: 4.20104129661425e-06\n",
      "[DEBUG]\tUpdated values for 600 steps. Difference between new and old table is: 2.8070459692774996e-08\n",
      "[DEBUG]\tUpdated values for 800 steps. Difference between new and old table is: 1.7462087331665543e-10\n",
      "[INFO]\tIn 2 iteration, current mean episode reward is 0.77.\n",
      "[DEBUG]\tUpdated values for 200 steps. Difference between new and old table is: 0.0004257477615745714\n",
      "[DEBUG]\tUpdated values for 400 steps. Difference between new and old table is: 1.4125290733302265e-05\n",
      "[DEBUG]\tUpdated values for 600 steps. Difference between new and old table is: 3.971402302571647e-07\n",
      "[DEBUG]\tUpdated values for 800 steps. Difference between new and old table is: 1.0301546324309463e-08\n",
      "[DEBUG]\tUpdated values for 1000 steps. Difference between new and old table is: 2.548876110175513e-10\n",
      "[INFO]\tIn 3 iteration, current mean episode reward is 0.688.\n",
      "[DEBUG]\tUpdated values for 200 steps. Difference between new and old table is: 0.00019132880435357436\n",
      "[DEBUG]\tUpdated values for 400 steps. Difference between new and old table is: 1.8012092146968417e-05\n",
      "[DEBUG]\tUpdated values for 600 steps. Difference between new and old table is: 1.3738229058951612e-06\n",
      "[DEBUG]\tUpdated values for 800 steps. Difference between new and old table is: 9.513477411404736e-08\n",
      "[DEBUG]\tUpdated values for 1000 steps. Difference between new and old table is: 6.230829824316331e-09\n",
      "[DEBUG]\tUpdated values for 1200 steps. Difference between new and old table is: 3.9353281744425317e-10\n",
      "[INFO]\tIn 4 iteration, current mean episode reward is 0.829.\n",
      "[INFO]\tIn 5 iteration, current mean episode reward is 0.867.\n",
      "We found policy is not changed anymore at itertaion 6. Current mean episode reward is 0.867. Stop training.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "# It may be confusing to call a trainer agent. But that's what we normally do.\n",
    "pi_agent = policy_iteration()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "print(\"Your policy iteration agent achieve {} mean episode reward. The optimal score \"\n",
    "      \"should be closed to {}.\".format(pi_agent.evaluate(), 0.86))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "pi_agent.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+-----+-----State Value Mapping-----+-----+-----+\n",
      "|     |   0 |   1 |   2 |   3 |   4 |   5 |   6 |   7 |\n",
      "|-----+-----+-----+-----+-----+-----+-----+-----+-----|\n",
      "| 0   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 1   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 2   |1.000|0.978|0.926|0.000|0.857|0.946|0.982|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 3   |1.000|0.935|0.801|0.475|0.624|0.000|0.945|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 4   |1.000|0.826|0.542|0.000|0.539|0.611|0.852|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 5   |1.000|0.000|0.000|0.168|0.383|0.442|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 6   |1.000|0.000|0.195|0.121|0.000|0.332|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 7   |1.000|0.732|0.463|0.000|0.277|0.555|0.777|0.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "pi_agent.print_table()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations! You have successfully implemented the policy iteration trainer (if and only if no error happens at the above cells). \n",
    "\n",
    "Here are few further problems for you to investigate:\n",
    "\n",
    "1. What is the impact of the discount factor gamma?\n",
    "2. What is the impact of the value function convergence criterion epsilon?\n",
    "\n",
    "If you are interested in doing more investigation (not limited to these two), feel free to open new cells at the end of this notebook and left a clear trace of your thinking and coding, which leads to extra credit if you do a good job. It's an optional job, and you can ignore it.\n",
    "\n",
    "Now let's continue our journey!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2.2: Value Iteration\n",
    "\n",
    "Recall the idea of value iteration. We update the state value: \n",
    "\n",
    "$$v_{k+1}(s) = \\max_a E_{s'} [r(s, a) + \\gamma v_{k}(s')]$$\n",
    "\n",
    "wherein the $s'$ is next state, $r$ is the reward, $v_{k}(s')$ is the next state value given by the old (not updated yet) value function. The expectation is computed among all possible transitions (given a state and action pair, it is possible to have many different next states, since the environment is not deterministic).\n",
    "\n",
    "The value iteration algorithm does not require an inner loop. It computes the expected return of all possible actions at a given state and uses the maximum of them as the state value. You can imagine it \"pretends\" we already have the optimal policy and run policy iteration based on it. Therefore we do not need to maintain a policy object in a trainer. We only need to retrieve the optimal policy using the same rule as policy iteration, given current value function.\n",
    "\n",
    "You should implement the trainer following the framework we already wrote for you. Please carefully go through the code and finish all `TODO` in it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "\n",
    "class ValueIterationTrainer(PolicyItertaionTrainer):\n",
    "    \"\"\"Note that we inherate Policy Iteration Trainer, to resue the\n",
    "    code of update_policy(). It's same since it get optimal policy from\n",
    "    current state-value table (self.table).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, gamma=1.0, env_name='FrozenLake8x8-v1'):\n",
    "        super(ValueIterationTrainer, self).__init__(gamma, None, env_name)\n",
    "\n",
    "    def train(self):\n",
    "        \"\"\"Conduct one iteration of learning.\"\"\"\n",
    "        # [TODO] value function may be need to be reset to zeros.\n",
    "        # if you think it should, than do it. If not, then move on.\n",
    "        # self.table = np.zeros((self.obs_dim,))\n",
    "\n",
    "        # In value iteration, we do not explicit require a\n",
    "        # policy instance to run. We update value function\n",
    "        # directly based on the transitions. Therefore, we\n",
    "        # don't need to run self.update_policy() in each step.\n",
    "        self.update_value_function()\n",
    "\n",
    "    def update_value_function(self):\n",
    "        old_table = self.table.copy()\n",
    "\n",
    "        for state in range(self.obs_dim):\n",
    "            state_value = 0\n",
    "\n",
    "            # [TODO] what should be de right state value?\n",
    "            # hint: try to compute the state_action_values first\n",
    "            act = self.policy(state)\n",
    "            transition_list = self._get_transitions(state, act)\n",
    "            for transition in transition_list:\n",
    "                prob = transition['prob']\n",
    "                reward = transition['reward']\n",
    "                next_state = transition['next_state']\n",
    "                done = transition['done']\n",
    "                state_value += prob * (reward + self.gamma * old_table[next_state])\n",
    "\n",
    "            self.table[state] = state_value\n",
    "\n",
    "        # Till now the one step value update is finished.\n",
    "        # You can see that we do not use a inner loop to update\n",
    "        # the value function like what we did in policy iteration.\n",
    "        # This is because to compute the state value, which is\n",
    "        # a expectation among all possible action given by a\n",
    "        # specified policy, we **pretend** already own the optimal\n",
    "        # policy (the max operation).\n",
    "\n",
    "    def evaluate(self):\n",
    "        \"\"\"Since in value itertaion we do not maintain a policy function,\n",
    "        so we need to retrieve it when we need it.\"\"\"\n",
    "        self.update_policy()\n",
    "        return super().evaluate()\n",
    "\n",
    "    def render(self):\n",
    "        \"\"\"Since in value itertaion we do not maintain a policy function,\n",
    "        so we need to retrieve it when we need it.\"\"\"\n",
    "        self.update_policy()\n",
    "        return super().render()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# Managing configurations of your experiments is important for your research.\n",
    "default_vi_config = dict(\n",
    "    max_iteration=10000,\n",
    "    evaluate_interval=100,  # don't need to update policy each iteration\n",
    "    gamma=1.0,\n",
    "    eps=1e-10\n",
    ")\n",
    "\n",
    "\n",
    "def value_iteration(train_config=None):\n",
    "    config = default_vi_config.copy()\n",
    "    if train_config is not None:\n",
    "        config.update(train_config)\n",
    "\n",
    "    # [TODO] initialize Value Iteration Trainer. Remember to pass\n",
    "    #  config['gamma'] to it.\n",
    "    trainer = ValueIterationTrainer(gamma=config['gamma'])\n",
    "\n",
    "    for i in range(config['max_iteration']):\n",
    "        old_state_value_table = trainer.table.copy()\n",
    "        old_policy_result = {\n",
    "            obs: trainer.policy(obs) for obs in range(trainer.obs_dim)\n",
    "        }\n",
    "        # train the agent\n",
    "        trainer.train()  # [TODO] please uncomment this line\n",
    "        # evaluate the result\n",
    "        if i % config['evaluate_interval'] == 0:\n",
    "            print(\"[INFO]\\tIn {} iteration, current \"\n",
    "                  \"mean episode reward is {}.\".format(\n",
    "                i, trainer.evaluate()\n",
    "            ))\n",
    "\n",
    "            # [TODO] compare the new policy with old policy to check should\n",
    "            #  we stop.\n",
    "            # [HINT] If new and old policy have same output given any\n",
    "            #  observation, them we consider the algorithm is converged and\n",
    "            #  should be stopped.\n",
    "            new_policy_result = {\n",
    "                obs: trainer.policy(obs) for obs in range(trainer.obs_dim)\n",
    "            }\n",
    "            should_stop = (new_policy_result == old_policy_result)\n",
    "            \n",
    "            if should_stop:\n",
    "                print(\"We found policy is not changed anymore at \"\n",
    "                      \"itertaion {}. Current mean episode reward \"\n",
    "                      \"is {}. Stop training.\".format(i, trainer.evaluate()))\n",
    "                break\n",
    "\n",
    "            if i > 3000:\n",
    "                print(\"You sure your codes is OK? It shouldn't take so many \"\n",
    "                      \"({}) iterations to train a policy iteration \"\n",
    "                      \"agent.\".format(\n",
    "                    i))\n",
    "\n",
    "    assert trainer.evaluate() > 0.8, \\\n",
    "        \"We expect to get the mean episode reward greater than 0.8. \" \\\n",
    "        \"But you get: {}. Please check your codes.\".format(trainer.evaluate())\n",
    "\n",
    "    return trainer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO]\tIn 0 iteration, current mean episode reward is 0.0.\n",
      "[INFO]\tIn 100 iteration, current mean episode reward is 0.892.\n",
      "[INFO]\tIn 200 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 300 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 400 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 500 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 600 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 700 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 800 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 900 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1000 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1100 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1200 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1300 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1400 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1500 iteration, current mean episode reward is 0.867.\n",
      "[INFO]\tIn 1600 iteration, current mean episode reward is 0.867.\n",
      "We found policy is not changed anymore at itertaion 1600. Current mean episode reward is 0.867. Stop training.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your value iteration agent achieve 0.867 mean episode reward. The optimal score should be almost 0.86.\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "print(\"Your value iteration agent achieve {} mean episode reward. The optimal score \"\n",
    "      \"should be almost {}.\".format(vi_agent.evaluate(), 0.86))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "vi_agent.render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+-----+-----State Value Mapping-----+-----+-----+\n",
      "|     |   0 |   1 |   2 |   3 |   4 |   5 |   6 |   7 |\n",
      "|-----+-----+-----+-----+-----+-----+-----+-----+-----|\n",
      "| 0   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 1   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 2   |1.000|0.978|0.926|0.000|0.857|0.946|0.982|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 3   |1.000|0.935|0.801|0.475|0.624|0.000|0.945|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 4   |1.000|0.826|0.542|0.000|0.539|0.611|0.852|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 5   |1.000|0.000|0.000|0.168|0.383|0.442|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 6   |1.000|0.000|0.195|0.121|0.000|0.332|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 7   |1.000|0.732|0.463|0.000|0.277|0.555|0.777|0.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Run this cell without modification\n",
    "\n",
    "vi_agent.print_table()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulation! You have successfully implemented the value iteration trainer (if and only if no error happens at the above cells). Few further problems for you to investigate:\n",
    "\n",
    "1. Do you see that some iteration during training yields better rewards than the final one?  Why does that happen?\n",
    "2. What is the impact of the discount factor gamma?\n",
    "3. What is the impact of the value function convergence criterion epsilon?\n",
    "\n",
    "If you are interested in doing more investigation (not limited to these two), feel free to open new cells at the end of this notebook and left a clear trace of your thinking and coding, which leads to extra credit if you do a good job. It's an optional job, and you can ignore it.\n",
    "\n",
    "Now let's continue our journey!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2.3: Compare two model-based agents\n",
    "\n",
    "Now we have two agents: `pi_agent` and `vi_agent`. They are believed to be the optimal policy in this environment. Can you compare the policy of two of them and use a clean and clear description or figures to show your conclusion?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+-----+-----State Value Mapping-----+-----+-----+\n",
      "|     |   0 |   1 |   2 |   3 |   4 |   5 |   6 |   7 |\n",
      "|-----+-----+-----+-----+-----+-----+-----+-----+-----|\n",
      "| 0   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 1   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 2   |1.000|0.978|0.926|0.000|0.857|0.946|0.982|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 3   |1.000|0.935|0.801|0.475|0.624|0.000|0.945|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 4   |1.000|0.826|0.542|0.000|0.539|0.611|0.852|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 5   |1.000|0.000|0.000|0.168|0.383|0.442|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 6   |1.000|0.000|0.195|0.121|0.000|0.332|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 7   |1.000|0.732|0.463|0.000|0.277|0.555|0.777|0.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "\n",
      "+-----+-----+-----State Value Mapping-----+-----+-----+\n",
      "|     |   0 |   1 |   2 |   3 |   4 |   5 |   6 |   7 |\n",
      "|-----+-----+-----+-----+-----+-----+-----+-----+-----|\n",
      "| 0   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 1   |1.000|1.000|1.000|1.000|1.000|1.000|1.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 2   |1.000|0.978|0.926|0.000|0.857|0.946|0.982|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 3   |1.000|0.935|0.801|0.475|0.624|0.000|0.945|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 4   |1.000|0.826|0.542|0.000|0.539|0.611|0.852|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 5   |1.000|0.000|0.000|0.168|0.383|0.442|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 6   |1.000|0.000|0.195|0.121|0.000|0.332|0.000|1.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "| 7   |1.000|0.732|0.463|0.000|0.277|0.555|0.777|0.000|\n",
      "|     |     |     |     |     |     |     |     |     |\n",
      "+-----+-----+-----+-----+-----+-----+-----+-----+-----+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Solve the TODO and remove `pass`\n",
    "\n",
    "# [TODO] try to compare two trained agents' policies\n",
    "# hint: trainer.print_table() may give you a better sense.\n",
    "vi_agent.print_table()\n",
    "pi_agent.print_table()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  (Right)\n",
      "SFFFFFFF\n",
      "FFFFFFFF\n",
      "FFFHFFFF\n",
      "FFFFFHFF\n",
      "FFFHFFFF\n",
      "FHHFFFHF\n",
      "FHFFHFHF\n",
      "FFFHFFF\u001B[41mG\u001B[0m\n",
      "Current observation: 63\n",
      "Current reward: 1.0\n",
      "Whether we are done: True\n",
      "info: {'prob': 0.3333333333333333}\n"
     ]
    }
   ],
   "source": [
    "# You can do more inverstigation here if you wish. Leave it blank if you don't.\n",
    "# Solve the TODOs and remove `pass`\n",
    "\n",
    "# Run 1000 steps for test, terminate if done.\n",
    "# You can run this cell multiples times.\n",
    "obs = env.reset()\n",
    "\n",
    "while True:\n",
    "    # take agent action\n",
    "    # [TODO] Uncomment next line\n",
    "    obs, reward, done, info = env.step(vi_agent.policy(obs))\n",
    "\n",
    "    # render the environment\n",
    "    env.render()  # [TODO] Uncomment this line\n",
    "\n",
    "    print(\"Current observation: {}\\nCurrent reward: {}\\n\"\n",
    "          \"Whether we are done: {}\\ninfo: {}\".format(\n",
    "        obs, reward, done, info\n",
    "    ))\n",
    "    wait(sleep=0.5)\n",
    "\n",
    "    # [TODO] terminate the loop if done\n",
    "    if done:\n",
    "        env.reset()\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "------\n",
    "\n",
    "## Conclusion and Discussion\n",
    "\n",
    "In this assignment, we learn how to use Gym package, how to use Object Oriented Programming idea to build a basic tabular RL algorithm.\n",
    "\n",
    "It's OK to leave the following cells empty. In the next markdown cell, you can write whatever you like. Like the suggestion on the course, the confusing problems in the assignments, and so on.\n",
    "\n",
    "Following the submission instruction in the assignment to submit your assignment to our staff. Thank you!\n",
    "\n",
    "------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "..."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}